﻿namespace CSP

module Extended =

    open CSP.Domain
    open CSP.Problem
    open CSP.Solving

    type private 'a extended =
        | Simple of 'a
        | Extended of 'a extended * 'a extended

    let private simple ext =
        match ext with
            | Simple x -> x
            | _ -> failwithf "Not a simple value: %A." ext

    let private extended ext =
        match ext with
            | Extended (l, r) -> l, r
            | _ -> failwithf "Not an extended value: %A." ext

    let private extendDomain domain = domain |> values |> List.map Simple |> ofList
    
    let private extendProblem problem =
        let ds = problem |> domains |> Map.toList |> List.map (fun (k, d) -> Simple k, extendDomain d) |> Map.ofList
        let cs = problem |> constraints |> List.map (fun (l,c,r) -> (Simple l), (fun x y -> c (simple x) (simple y)), (Simple r)) 
        create ds cs

    let private combineDomains key1 key2 problem =
        let key = Extended (key1,key2)
        let domain =
            let d1, d2 = problem |> domains |> fun m -> Map.find key1 m, Map.find key2 m
            [ for x in values d1 do for y in values d2 -> Extended (x,y) ] |> ofList
        let p = addDomain key domain problem
        let c1 = fun x y -> let xl, _ = extended x in xl = y
        let c2 = fun x y -> let _, xr = extended x in xr = y
    
        problem |> addDomain key domain |> addConstraint key c1 key1 |> addConstraint key c2 key2  

    let rec private ensureKey key problem =
        match key with
            | Simple k -> problem
            | Extended (k1,k2) -> if problem |> domains |> Map.containsKey key then problem
                                  else problem |> ensureKey k1 |> ensureKey k2 |> combineDomains k1 k2

    let rec private toExtended list =
        match list with
            | [] -> failwith "Empty list."
            | [x] -> Simple x
            | _ ->
                let count = List.length list
                let length = count / 2
                let s l = List.map snd l
                let l, r = list |> List.mapi (fun i x -> (i,x)) |> List.partition (fun (i,_) -> i < length) |> fun (x,y) -> (s x, s y)
                Extended (toExtended l, toExtended r)

    let rec private fromExtended ext =
        match ext with
            | Simple x -> [x]
            | Extended (x,y) -> List.append (fromExtended x) (fromExtended y)

    let private addNullary value (problem: problem<'k,'a>) =
        if value then problem
        else let emptyDomain: 'a domain = ofList []
             let domains = problem |> domains |> Map.map (fun k _ -> emptyDomain)
             create domains []

    let private addUnary key pred problem =
        let domain = problem |> domains |> Map.find key |> filter (fun x -> pred x)
        let domains = problem |> domains |> Map.add key domain
        create domains (constraints problem)

    let private add012Constraint (keys,pred) problem =
        match keys with
            | [] -> addNullary (pred []) problem
            | [k] -> addUnary k (fun x -> pred [x]) problem
            | [k1;k2] -> addConstraint k1 (fun x y -> pred [x;y]) k2 problem
            | _ -> failwithf "Expected constraint with arity 0, 1 or 2, but got %A." (keys,pred)

    let private addExtendedConstraint keys pred (problem: problem<'k extended,'a extended>) =
        match keys with
            | [] -> addNullary (pred []) problem
            | [key] -> addUnary (Simple key) (fun x -> pred [simple x]) problem
            | [k1;k2] ->
                addConstraint (Simple k1) (fun x y -> pred [simple x; simple y]) (Simple k2) problem
            | _ ->
                let k1,k2 = keys |> toExtended |> extended
                let c = fun x y -> (x,y) |> Extended |> fromExtended |> pred
                problem |> ensureKey k1 |> ensureKey k2 |> addConstraint k1 c k2

    /// Creates a nullary constraint.
    let nullary b = [], fun [] -> b

    /// Creates a unary constraint.
    let unary k p = [k], fun [x] -> p x

    /// Creates a binary constraint.
    let binary k1 k2 p = [k1;k2], fun [x;y] -> p x y

    /// Creates a ternary constraint.
    let ternary k1 k2 k3 p = [k1;k2;k3], fun [x;y;z] -> p x y z

    /// Creates a quaternary constraint.
    let quaternary k1 k2 k3 k4 p = [k1;k2;k3;k4], fun [x;y;z;u] -> p x y z u

    let private arity (keys,_) = List.length keys

    /// Finds all solutions of a general CSP with global checks.
    let solve (st: 's) (update: 'k -> 'a -> 's -> 's) (check: 's -> bool) (domains: Map<'k,'a domain>) (constraints: #seq<_>) =
        let c012, c3 = constraints |> List.ofSeq |> List.partition (fun c -> arity c <= 2)
        let p = create domains []
        let p = c012 |> Seq.fold (fun p c -> add012Constraint c p) p
        if List.isEmpty c3 then solve st update check p
        else
            let p = extendProblem p
            let p = c3 |> Seq.fold (fun p (keys,pred) -> addExtendedConstraint keys pred p) p
            let t s = domains |> Map.map (fun k _ -> simple <| Map.find (Simple k) s)
            let update' k x s = match k with
                                | Simple k' -> update k' (simple x) s
                                | _ -> s
            solve st update' check p |> Seq.map t

    /// Finds one solution of a general CSP with global checks.
    let find st update check domains constraints = solve st update check domains constraints |> Utils.SeqUtils.safeHead

    /// Finds all solutions of a general CSP without global checks.
    let simpleSolve domains constraints = solve simpleState simpleUpdate simpleCheck domains constraints

    /// Finds one solution of a general CSP without global checks.
    let simpleFind domains constraints = find simpleState simpleUpdate simpleCheck domains constraints